package org.finesys.common.core.support.concurrent;

import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Future;
import java.util.concurrent.RejectedExecutionException;

public class TaskManager {

    private static final ConcurrentHashMap<String, List<Future<?>>> TASK_MAP = new ConcurrentHashMap<>();


    /**
     * 提交任务到异步线程池
     *
     * @param id   任务唯一标识
     * @param task 要提交的任务
     * @throws IllegalArgumentException 如果 id 或 task 为空，则抛出此异常
     * @throws RuntimeException         如果提交任务到线程池失败，则抛出此异常
     */
    public static void submitTask(String id, Callable<?> task) {
        if (id == null || task == null) {
            throw new IllegalArgumentException("id and task cannot be null");
        }
        try {
            Future<?> future = AsyncThreadPool.getThreadPool().submit(task);
            List<Future<?>> futures = TASK_MAP.computeIfAbsent(id, k -> new java.util.ArrayList<>());
            futures.add(future);
        } catch (RejectedExecutionException e) {
            throw new RuntimeException("Failed to submit task to thread pool", e);
        }
    }

    /**
     * 从任务映射中移除指定ID的任务
     *
     * @param id 需要移除的任务ID
     * @return 被移除的任务对象，如果不存在则返回null
     */
    public List<Future<?>> removeTask(String id) {
        if (id == null) {
            return null;
        }
        synchronized (TASK_MAP) {
            return TASK_MAP.remove(id);
        }
    }

    /**
     * 根据给定的任务ID获取该任务对应的Future对象数量。
     *
     * @param id 任务ID
     * @return 该任务对应的Future对象数量，如果任务不存在则返回0
     */
    public int getCount(String id) {
        if (TASK_MAP.containsKey(id)) {
            List<Future<?>> futures = TASK_MAP.get(id);
            return futures == null ? 0 : futures.size();
        }
        return 0;
    }

}